# 自定义损失函数

import torch
import torch.nn as nn

class PathLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.seq_loss = nn.CrossEntropyLoss()
        self.img_loss = nn.MSELoss()
        
        # 修改为更合理的图像重建网络
        self.img_decoder = nn.Sequential(
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 1024),
            nn.ReLU(),
            nn.Linear(1024, 128*128*3),  # 直接重建整图
            nn.Tanh()  # 输出在[-1,1]范围
        )

    def forward(self, outputs, targets):
        # 序列损失
        seq_loss = self.seq_loss(
            outputs['seq'].flatten(0, 1),
            targets['tokens'].flatten()
        )
        
        # 图像重建损失（全图级别）
        img_pred = self.img_decoder(outputs['img'].mean(dim=1))
        img_pred = img_pred.view(-1, 3, 128, 128)
        img_loss = self.img_loss(img_pred, targets['image'])
        
        return {
            'seq': seq_loss,
            'img': img_loss,
            'total': seq_loss + img_loss
        }

class GANLoss:
    @staticmethod
    def adv_loss(logits, target_is_real=True):
        target = torch.ones_like(logits) if target_is_real else torch.zeros_like(logits)
        return nn.BCEWithLogitsLoss()(logits, target)